# Code for Neural Interaction Unit (NIU).
# Reference:
# Zhang, J.; Zhou, Y.; Lei, Y.-K.; Yang, Y. I.; Gao, Y. Q.,
# Molecular CT: unifying geometry and representation learning for molecules at different scales [J/OL].
# arXiv preprint, 2020: arXiv:2012.11816 [2020-12-22]. https://arxiv.org/abs/2012.11816

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Tuple, Callable
from jax import nn as nn
from ...utils.base import PositionalEmbedding, MultiheadAttention
from ...utils.base import HyperformerPairBlock
from ...utils.activation import get_activation

class InteractionUnit(hk.Module):
    r"""Interaction unit for MolCT.

    ### Args:

        dim_feature (int): Dimension of node and edge vectors.

        n_heads (int): Number of heads in multi-head attention.

        dim_outer_pdct (int): Dimension of outer product. Default: 32.
 
        num_transition (int): Number of transition in hyperformer. Default: 4.

        is_edge_update (bool): Whether to update edge vectors. Default: False.

        activation (Callable, str): Activation function, Default: 'silu'.

        fp_type (dtype): Floating point type. Default: jnp.float32.

        name (str): Name of the module. Default: 'neural_interaction_unit'.
    
    """

    def __init__(self,
                 dim_feature: int,
                 n_heads: int,
                 dim_outer_pdct: int = 32,
                 num_transition: int = 4,
                 is_edge_update: bool = False,
                 activation: Optional[Union[str, Callable]] = 'silu',
                 fp_type = jnp.float32,
                 name: str = "neural_interaction_unit",
                 ):
        super().__init__(name=name)

        self.fp_type = fp_type
        self.is_edge_update = is_edge_update
        self.activation = get_activation(activation) # not used in public molct

        self.positional_embedding = PositionalEmbedding(dim_feature=dim_feature, fp_type=fp_type) # type: ignore
        self.multi_head_attention = MultiheadAttention(dim_feature=dim_feature, n_heads=n_heads) # type: ignore

        self.hyperformer_pair_block = None
        if self.is_edge_update:
            self.hyperformer_pair_block = HyperformerPairBlock(dim_feature, dim_outer_pdct, num_transition, fp_type) # type: ignore

        self.node_norm = hk.LayerNorm(axis=-1, 
                                      create_scale=True, 
                                      create_offset=True, 
                                      name="node_norm") # type: ignore
        self.edge_norm = hk.LayerNorm(axis=-1, 
                                      create_scale=True, 
                                      create_offset=True, 
                                      name="edge_norm") # type: ignore
        
    def __call__(self,
                 node_vec: jnp.ndarray,
                 node_mask: jnp.ndarray,
                 edge_vec: jnp.ndarray,
                 edge_mask: jnp.ndarray,
                 edge_cutoff: jnp.ndarray,
                 ):
        r"""
        ### Args:

            node_vec (jnp.ndarray): Node vectors. Shape: (A, F).

            node_mask (jnp.ndarray): Node masks. Shape: (A,).

            edge_vec (jnp.ndarray): Edge vectors. Shape: (A, A, F).

            edge_mask (jnp.ndarray): Edge masks. Shape: (A, A).

            edge_cutoff (jnp.ndarray): Edge cutoffs. Shape: (A, A).
        """
        
        # Initializing the encoders for nodes and edges.
        def _node_encoder(node_vec: jnp.ndarray,
                          edge_vec: jnp.ndarray,
                          edge_mask: jnp.ndarray,
                          edge_cutoff: jnp.ndarray,
                          time_signal: float = 0.0,
                          ) -> jnp.ndarray:
            
            query, key, value = self.positional_embedding(node_vec,
                                                          edge_vec,
                                                          time_signal,)
            delta_node_vec = self.multi_head_attention(query,
                                                       key,
                                                       value,
                                                       edge_mask,
                                                       edge_cutoff,)
            node_new = node_vec + delta_node_vec.squeeze(axis=1)
            node_new = self.node_norm(node_new)

            return node_new

        def _edge_encoder(node_vec: jnp.ndarray,
                          edge_vec: jnp.ndarray,
                          node_mask: jnp.ndarray,
                          edge_mask: jnp.ndarray,
                          ) -> jnp.ndarray:
            
            edge_vec = self.hyperformer_pair_block(node_vec=node_vec,
                                                   edge_vec=edge_vec,
                                                   node_mask=node_mask,
                                                   edge_mask=edge_mask,) # type: ignore
            edge_vec = self.edge_norm(edge_vec)

            return edge_vec

        if self.is_edge_update:
            node_new = _node_encoder(
                        node_vec=node_vec,
                        edge_vec=edge_vec,
                        edge_mask=edge_mask,
                        edge_cutoff=edge_cutoff,)

            edge_new = _edge_encoder(
                        node_vec=node_vec,
                        edge_vec=edge_vec,
                        node_mask=node_mask,
                        edge_mask=edge_mask,)
        else:
            node_new = _node_encoder(
                        node_vec=node_vec,
                        edge_vec=edge_vec,
                        edge_mask=edge_mask,
                        edge_cutoff=edge_cutoff,)

            edge_new = edge_vec

        return node_new, edge_new                                   
            

